import torch
import torch.nn as nn
from collections import OrderedDict

class InceptionV3(nn.Module):
    def __init__(self):
        super().__init__()

    def __create_basic_nn(self):
        basic_nn = nn.Sequential()
        basic_nn.add_module("conv1", ConvBlock(3, 32, 3, 2))
        basic_nn.add_module("conv2", ConvBlock(32, 32, 3, 1))
        basic_nn.add_module("conv3", ConvBlock(32, 64, 3, 1, 1))
        basic_nn.add_module("pool", nn.MaxPool2d(3, 2))
        basic_nn.add_module("conv4", ConvBlock(64, 80, 3, 1))
        basic_nn.add_module("conv5", ConvBlock(80, 192, 3, 2))
        basic_nn.add_module("conv5", ConvBlock(192, 288, 3, 1, 1))


class ConvBlock(nn.Module):
    def __init__(self, in_c, out_c, kernel_size, stride, padding=0):
        super().__init__()
        self.conv = nn.Sequential(OrderedDict([
            ("conv", nn.Conv2d(in_c, out_c, kernel_size=kernel_size, stride=stride, padding=padding, bias=False)),
            ("bn", nn.BatchNorm2d(out_c)),
            ("relu", nn.ReLU(True))
        ]))

    def forward(self, x):
        return self.conv(x)

class Inception3x3(nn.Module):
    def __init__(self, in_c, out_c):
        super().__init__()
        self.branch_1 = nn.Sequential(OrderedDict([
            ("conv1x1", ConvBlock(in_c, 64, 1, 1)),
            ("conv3x3_1", ConvBlock(64, out_c, 3, 1)),
            ("conv3x3_2", ConvBlock(out_c, out_c, 3, 1))
        ]))

        self.branch_2 = nn.Sequential(OrderedDict([
            ("conv1x1", ConvBlock(in_c, 48, 1, 1)),
            ("conv3x3", ConvBlock(48, 64, 3, 1))
        ]))

        self.branch_3 = nn.Sequential(OrderedDict([
            ("conv1x1", nn.AvgPool2d(3, 1, 1)),
            ("conv3x3", ConvBlock(48, 64, 3, 1))
        ]))

        self.branch_4 = nn.Sequential(OrderedDict([
            ("conv1x1", ConvBlock(48, 64, 3, 1))
        ]))

    def forward(self, x):
        out_branch_1 = self.branch_1(x)
        out_branch_2 = self.branch_2(x)
        out_branch_3 = self.branch_3(x)
        out_branch_4 = self.branch_4(x)

        return torch.cat((out_branch_1, out_branch_2, out_branch_3, out_branch_4), dim=1)
